#!/usr/bin/python
# -*-coding:utf-8-*-
import os
import pandas as pd
import numpy as np

from base.KeyTransform import stock_code_map
from base.KeyTransform import trade_date_map

from base.GPModelSave import load_gpmodel
from base.GPModelSave import model_predict

### 底层读取数据的依赖（不提供）
from zbc_factor_lib.base.factors_library_base import NewRQFactorLib as DataReader

db = 'gp_factors'
data_reader = DataReader(db=db)

# TODO - 固定的参数
label_data_dir = './zbc_gplearn_factor_mining/label_data'
label_data_filename = 'financial_label_data'

test_data_dir = './zbc_gplearn_factor_mining/raw_X_data'
test_data_filename = 'X_financial_statement_data_v1'

# TODO - 保存模型结果
save_model_dir = './zbc_gplearn_factor_mining/model/financial_factors_v1'

def load_data_v1(selected_X_list=None, **kwargs):
    # TODO - 读取数据
    # TODO - 测试
    if selected_X_list is None:
        test_X_data = pd.read_hdf(os.path.join(test_data_dir, test_data_filename + '.h5'))
    else:
        test_X_data = pd.read_hdf(os.path.join(test_data_dir, test_data_filename + '.h5'),
                                  columns=selected_X_list)

    # TODO - Map Keys
    test_X_data = test_X_data.reset_index()

    test_X_data = trade_date_map(test_X_data, key='report_date', reverse=False)
    test_X_data = stock_code_map(test_X_data, key='stock_code', reverse=False)

    test_X_data = test_X_data.set_index(['report_date', 'stock_code'])

    # TODO - get label
    concat_label_data = pd.read_hdf(os.path.join(label_data_dir, label_data_filename+'.h5'))

    neu_keys = concat_label_data.loc[test_X_data.index, ['ci1_code', 'cap']].copy()
    metric_features_keys = test_X_data.index

    del concat_label_data
    # del test_X_data

    return test_X_data, \
            metric_features_keys, \
            neu_keys

def get_factor_main(model_filname, start_date=None, end_date=None, raw_data=None, keys=None, neu_keys=None):
    # model_filname = 'zg02_ths_user_hehavior_factor_mining_v2_v2'
    model = load_gpmodel(os.path.join(save_model_dir, model_filname+'.pkl'))

    Xnames = [k for k in model if isinstance(k, str)]

    # TODO - 去掉重复的X
    Xnames = list(set(Xnames))

    if raw_data is None:
        used_raw_data, keys, neu_keys = load_data_v1(Xnames)
    else:
        used_raw_data = raw_data[Xnames]

    # [p.name if not isinstance(p, (str, int, float)) else p for p in model]

    factor_data = model_predict(X=used_raw_data,
                                program=[model, 5],
                                keys=keys,
                                neu_keys=neu_keys,
                                by_name=True,
                                key_by_path=False)

    factor_data = pd.DataFrame(factor_data, index=keys)
    factor_data.columns = [model_filname]

    factor_data = factor_data.reset_index()

    factor_data = stock_code_map(factor_data, key='stock_code', reverse=True)
    factor_data = trade_date_map(factor_data, key='date', reverse=True)

    if start_date is None and end_date is not None:
        factor_data = factor_data[factor_data['date'] <= pd.to_datetime(end_date)]
    elif start_date is not None and end_date is None:
        factor_data = factor_data[factor_data['date'] >= pd.to_datetime(start_date)]
    elif start_date is not None and end_date is not None:
        factor_data = factor_data[(factor_data['date'] >= pd.to_datetime(start_date)) &
                                  (factor_data['date'] <= pd.to_datetime(end_date))]

    return factor_data.reset_index(drop=True)

if __name__ == '__main__':
    start_date = '2017-01-01'
    end_date = '2020-03-31'

    raw_data, keys, neu_keys = load_data_v1()
    print('load init data done!')

    model_filname_list = [
        'zg02_financial_factor_mining_v1_v%s' % id for id in range(37, 42)
    ]

    for model_filname in model_filname_list:
        factor_data = get_factor_main(model_filname,
                                      start_date=start_date,
                                      end_date=end_date,
                                      raw_data=raw_data,
                                      keys=keys,
                                      neu_keys=neu_keys)

        data_reader.update_data(data=factor_data)

        data_reader.create_factor_table(filename=model_filname, main_columns=['date', 'stock_code'])

        print(model_filname, 'created')

